import os
import torch
import torch.nn as nn
import torch.distributed as dist
import numpy as np

from pcdet.models.model_utils import model_nms_utils
from pcdet.utils.common_utils import create_logger
from pcdet.ops.iou3d_nms import iou3d_nms_utils
from pcdet.models import backbones_2d, backbones_3d_stereo, backbones_3d_lidar,dense_heads
from pcdet.models.backbones_2d import map_to_bev
from pcdet.models.backbones_3d_lidar import vfe
from pcdet.models.dense_heads.depth_loss_head import DepthLossHead
from pcdet.models.dense_heads.voxel_loss_head import VoxelLossHead


class StereoDetector3DTemplate(nn.Module):
    def __init__(self, model_cfg, num_class, dataset):
        super().__init__()
        self.model_cfg = model_cfg
        self.num_class = num_class
        self.dataset = dataset
        self.class_names = dataset.class_names
        self.register_buffer('global_step', torch.LongTensor(1).zero_())

        self.module_topology = [
            'vfe','backbone_3d', 'backbone_lidar', 'map_to_bev_module',
            'backbone_2d', 'dense_head_2d', 'dense_head', 
            'depth_loss_head', 'voxel_loss_head', 
        ]
        if 'DENSE_HEAD' in self.model_cfg and getattr(self.model_cfg.DENSE_HEAD, 'do_feature_imitation', False):
            if hasattr(self.model_cfg, 'LIDAR_MODEL'):
                self.module_topology.insert(0, 'lidar_model')
            else:
                self.module_topology.insert(0, 'stereo_model')

    def train(self, mode=True):
        self.training = mode
        for module in self.children():
            if module in self.model_info_dict['fixed_module_list']:
                module.eval()
            else:
                module.train(mode)
        return self

    @property
    def mode(self):
        return 'TRAIN' if self.training else 'TEST'

    def update_global_step(self):
        self.global_step += 1

    def build_networks(self, trt=False):
        model_info_dict = {
            'module_list': [],
            'fixed_module_list': [],
            'grid_size': getattr(self.dataset, 'stereo_grid_size', self.dataset.grid_size),
            'point_cloud_range': self.dataset.point_cloud_range,
            'voxel_size': getattr(self.dataset, 'stereo_voxel_size', self.dataset.voxel_size),
            'boxes_gt_in_cam2_view': getattr(self.dataset, 'boxes_gt_in_cam2_view', False),
            'stereo_point_cloud_range': self.dataset.stereo_point_cloud_range,
            'max_crop_shape': getattr(self.dataset, 'max_crop_shape', None),
            'num_rawpoint_features': self.dataset.point_feature_encoder.num_point_features,
            'num_point_features': self.dataset.point_feature_encoder.num_point_features,
            'trt':trt,
        }

        for module_name in self.module_topology:
            module, model_info_dict = getattr(self, 'build_%s' % module_name)(
                model_info_dict=model_info_dict
            )
            self.add_module(module_name, module)
        self.model_info_dict = model_info_dict

        if self.model_cfg.get('FREEZE_RPN', False):
            for module_name in self.module_topology:
                module = getattr(self, module_name, None)
                if module is None or module_name == 'roi_head': continue
                for param in module.parameters():
                    param.requires_grad_(False)

        return model_info_dict['module_list']

    def build_lidar_model(self, model_info_dict):
        from ..detectors_lidar import build_detector as build_lidar_detector

        if self.model_cfg.get('LIDAR_MODEL', None) is None:
            return None, model_info_dict

        lidar_model = build_lidar_detector(self.model_cfg.LIDAR_MODEL, self.num_class, self.dataset)
        for param in lidar_model.parameters():
            param.requires_grad_(False)
        model_info_dict['module_list'].append(lidar_model)
        model_info_dict['fixed_module_list'].append(lidar_model)
        model_info_dict['num_point_features'] = lidar_model.num_point_features
        model_info_dict['backbone_channels'] = lidar_model.backbone_channels \
            if hasattr(lidar_model, 'backbone_channels') else None
        
        try:
            logger = create_logger(rank=dist.get_rank())
        except:
            logger = create_logger()
        if hasattr(self.model_cfg.LIDAR_MODEL, 'PRETRAINED_MODEL') and self.model_cfg.LIDAR_MODEL.PRETRAINED_MODEL:
            lidar_model.load_params_from_file(
                filename=self.model_cfg.LIDAR_MODEL.PRETRAINED_MODEL, to_cpu=True, logger=logger)

        return lidar_model, model_info_dict

    def build_stereo_model(self, model_info_dict):
        from ..detectors_stereo import build_detector as build_stereo_detector

        if self.model_cfg.get('STEREO_MODEL', None) is None:
            return None, model_info_dict

        stereo_model = build_stereo_detector(self.model_cfg.STEREO_MODEL, self.num_class, self.dataset)
        for param in stereo_model.parameters():
            param.requires_grad_(False)
        model_info_dict['module_list'].append(stereo_model)
        model_info_dict['fixed_module_list'].append(stereo_model)

        try:
            logger = create_logger(rank=dist.get_rank())
        except:
            logger = create_logger()
        if hasattr(self.model_cfg.STEREO_MODEL, 'PRETRAINED_MODEL') and self.model_cfg.STEREO_MODEL.PRETRAINED_MODEL:
            stereo_model.load_params_from_file(
                filename=self.model_cfg.STEREO_MODEL.PRETRAINED_MODEL, to_cpu=True, logger=logger)

        return stereo_model, model_info_dict
    
    def build_vfe(self, model_info_dict):
        if self.model_cfg.get('VFE', None) is None:
            return None, model_info_dict

        vfe_module = vfe.__all__[self.model_cfg.VFE.NAME](
            model_cfg=self.model_cfg.VFE,
            num_point_features=model_info_dict['num_rawpoint_features'],
            point_cloud_range=model_info_dict['point_cloud_range'],
            voxel_size=model_info_dict['voxel_size'],
            grid_size=model_info_dict['grid_size'],
            # depth_downsample_factor=model_info_dict['depth_downsample_factor']
        )
        model_info_dict['num_point_features'] = vfe_module.get_output_feature_dim()
        model_info_dict['module_list'].append(vfe_module)
        return vfe_module, model_info_dict
    
    def build_backbone_lidar(self, model_info_dict):
        if self.model_cfg.get('BACKBONE_LIDAR', None) is None:
            return None, model_info_dict
        backbone_lidar_module = backbones_3d_lidar.__all__[self.model_cfg.BACKBONE_LIDAR.NAME](
            model_cfg=self.model_cfg.BACKBONE_LIDAR,
            input_channels=model_info_dict['num_point_features'],
            grid_size=model_info_dict['grid_size'],
            voxel_size=model_info_dict['voxel_size'],
            point_cloud_range=model_info_dict['point_cloud_range']
        )
        model_info_dict['module_list'].append(backbone_lidar_module)
        model_info_dict['num_point_features'] = backbone_lidar_module.num_point_features
        model_info_dict['backbone_channels'] = backbone_lidar_module.backbone_channels \
            if hasattr(backbone_lidar_module, 'backbone_channels') else None
        return backbone_lidar_module, model_info_dict
    
    def build_backbone_3d(self, model_info_dict):
        if self.model_cfg.get('BACKBONE_3D', None) is None:
            return None, model_info_dict
        if 'trt' in model_info_dict and model_info_dict['trt']:
            self.model_cfg.BACKBONE_3D.NAME += 'TRT'

        backbone_3d_module = backbones_3d_stereo.__all__[self.model_cfg.BACKBONE_3D.NAME](
            model_cfg=self.model_cfg.BACKBONE_3D,
            class_names=self.class_names,
            grid_size=model_info_dict['grid_size'],
            voxel_size=model_info_dict['voxel_size'],
            point_cloud_range=model_info_dict['point_cloud_range'],
            boxes_gt_in_cam2_view=model_info_dict['boxes_gt_in_cam2_view'],
            stereo_point_cloud_range=model_info_dict['stereo_point_cloud_range'],
            max_crop_shape=model_info_dict['max_crop_shape'])
        model_info_dict['module_list'].append(backbone_3d_module)
        return backbone_3d_module, model_info_dict

    def build_map_to_bev_module(self, model_info_dict):
        if self.model_cfg.get('MAP_TO_BEV', None) is None:
            return None, model_info_dict

        map_to_bev_module = map_to_bev.__all__[self.model_cfg.MAP_TO_BEV.NAME](
            model_cfg=self.model_cfg.MAP_TO_BEV,
            grid_size=model_info_dict['grid_size'],
        )
        model_info_dict['module_list'].append(map_to_bev_module)
        model_info_dict['num_bev_features'] = map_to_bev_module.num_bev_features
        return map_to_bev_module, model_info_dict

    def build_backbone_2d(self, model_info_dict):
        if self.model_cfg.get('BACKBONE_2D', None) is None:
            return None, model_info_dict
        num_bev_features=model_info_dict.get('num_bev_features', None)
        backbone_2d_module = backbones_2d.__all__[self.model_cfg.BACKBONE_2D.NAME](
            model_cfg=self.model_cfg.BACKBONE_2D,
            input_channels=num_bev_features if num_bev_features else self.model_cfg.BACKBONE_2D.get('num_bev_features', None),
        )
        model_info_dict['module_list'].append(backbone_2d_module)
        model_info_dict['num_bev_features'] = backbone_2d_module.num_bev_features
        return backbone_2d_module, model_info_dict

    def build_dense_head_2d(self, model_info_dict):
        if self.model_cfg.get('DENSE_HEAD_2D', None) is None:
            return None, model_info_dict
        if self.model_cfg.DENSE_HEAD_2D.NAME == 'MMDet2DHead':
            dense_head_module = dense_heads.__all__[self.model_cfg.DENSE_HEAD_2D.NAME](
                model_cfg=self.model_cfg.DENSE_HEAD_2D
            )
            model_info_dict['module_list'].append(dense_head_module)
            return dense_head_module, model_info_dict
        else:
            dense_head_module = dense_heads.__all__[self.model_cfg.DENSE_HEAD_2D.NAME](
                model_cfg=self.model_cfg.DENSE_HEAD_2D,
                input_channels=32,
                num_class=self.num_class,
                class_names=self.class_names,
                grid_size=model_info_dict['grid_size'],
                point_cloud_range=model_info_dict['point_cloud_range'],
                predict_boxes_when_training=self.model_cfg.get('ROI_HEAD', False)
            )
            model_info_dict['module_list'].append(dense_head_module)
            return dense_head_module, model_info_dict

    def build_dense_head(self, model_info_dict):
        if self.model_cfg.get('DENSE_HEAD', None) is None:
            return None, model_info_dict
        dense_head_module = dense_heads.__all__[self.model_cfg.DENSE_HEAD.NAME](
            model_cfg=self.model_cfg.DENSE_HEAD,
            input_channels=model_info_dict['num_bev_features'],
            num_class=self.num_class if not self.model_cfg.DENSE_HEAD.CLASS_AGNOSTIC else 1,
            class_names=self.class_names,
            grid_size=model_info_dict['grid_size'],
            point_cloud_range=model_info_dict['point_cloud_range'],
            predict_boxes_when_training=self.model_cfg.get('ROI_HEAD', False) or self.model_cfg.DENSE_HEAD.get('predict_boxes_when_training', False),
            voxel_size=model_info_dict.get('voxel_size', False)
        )
        model_info_dict['module_list'].append(dense_head_module)
        return dense_head_module, model_info_dict

    def build_depth_loss_head(self, model_info_dict):
        # import pdb; pdb.set_trace()
        assert not self.model_cfg.BACKBONE_3D.get('front_surface_depth', False) or \
                (self.model_cfg.BACKBONE_3D.get('drop_psv_loss', False) or self.model_cfg.BACKBONE_3D.get('drop_psv', False))
        if (self.model_cfg.get('DEPTH_LOSS_HEAD', None) is None or self.model_cfg.BACKBONE_3D.get('drop_psv', False)) and \
            not self.model_cfg.BACKBONE_3D.get('front_surface_depth', False):
            return None, model_info_dict
        depth_loss_head = DepthLossHead(
            model_cfg=self.model_cfg.DEPTH_LOSS_HEAD,
            point_cloud_range=getattr(model_info_dict, 'stereo_point_cloud_range', model_info_dict['point_cloud_range']),
        )
        model_info_dict['module_list'].append(depth_loss_head)
        return depth_loss_head, model_info_dict
    
    def build_vfe(self, model_info_dict):
        if self.model_cfg.get('VFE', None) is None:
            return None, model_info_dict

        vfe_module = vfe.__all__[self.model_cfg.VFE.NAME](
            model_cfg=self.model_cfg.VFE,
            num_point_features=model_info_dict['num_rawpoint_features'],
            point_cloud_range=model_info_dict['point_cloud_range'],
            voxel_size=model_info_dict['voxel_size'],
            grid_size=model_info_dict['grid_size'],
            # depth_downsample_factor=model_info_dict['depth_downsample_factor']
        )
        model_info_dict['num_point_features'] = vfe_module.get_output_feature_dim()
        model_info_dict['module_list'].append(vfe_module)
        return vfe_module, model_info_dict
    
    def build_voxel_loss_head(self, model_info_dict):
        if self.model_cfg.get('VOXEL_LOSS_HEAD', None) is None:
            return None, model_info_dict
        voxel_loss_head = VoxelLossHead(
            model_cfg=self.model_cfg.VOXEL_LOSS_HEAD,
            point_cloud_range=getattr(model_info_dict, 'stereo_point_cloud_range', model_info_dict['point_cloud_range']),
        )
        model_info_dict['module_list'].append(voxel_loss_head)
        return voxel_loss_head, model_info_dict

    def forward(self, **kwargs):
        raise NotImplementedError

    def post_processing(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size:
                batch_cls_preds: (B, num_boxes, num_classes | 1) or (
                    N1+N2+..., num_classes | 1)
                                or [(B, num_boxes, num_class1), (B, num_boxes, num_class2) ...]
                multihead_label_mapping: [(num_class1), (num_class2), ...]
                batch_box_preds: (B, num_boxes, 7+C) or (N1+N2+..., 7+C)
                cls_preds_normalized: indicate whether batch_cls_preds is normalized
                batch_index: optional (N1+N2+...)
                has_class_labels: True/False
                roi_labels: (B, num_rois)  1 .. num_classes
                batch_pred_labels: (B, num_boxes, 1)
        Returns:

        """
        post_process_cfg = self.model_cfg.POST_PROCESSING
        batch_size = batch_dict['batch_size']
        recall_dict = {}
        pred_dicts = []
        for index in range(batch_size):
            if batch_dict.get('batch_index', None) is not None:
                assert batch_dict['batch_box_preds'].shape.__len__() == 2
                batch_mask = (batch_dict['batch_index'] == index)
            else:
                assert batch_dict['batch_box_preds'].shape.__len__() == 3
                batch_mask = index

            box_preds = batch_dict['batch_box_preds'][batch_mask]
            src_box_preds = box_preds

            if not isinstance(batch_dict['batch_cls_preds'], list):
                cls_preds = batch_dict['batch_cls_preds'][batch_mask]

                src_cls_preds = cls_preds
                assert cls_preds.shape[1] in [1, self.num_class]

                if not batch_dict['cls_preds_normalized']:
                    cls_preds = torch.sigmoid(cls_preds)
            else:
                cls_preds = [x[batch_mask]
                             for x in batch_dict['batch_cls_preds']]
                src_cls_preds = cls_preds
                if not batch_dict['cls_preds_normalized']:
                    cls_preds = [torch.sigmoid(x) for x in cls_preds]

            if post_process_cfg.NMS_CONFIG.MULTI_CLASSES_NMS:
                if not isinstance(cls_preds, list):
                    cls_preds = [cls_preds]
                    multihead_label_mapping = [torch.arange(1, self.num_class + 1, device=cls_preds[0].device)]
                else:
                    multihead_label_mapping = batch_dict['multihead_label_mapping']

                # NOTE: different with OpenPCDet
                if batch_dict.get('has_class_labels', False):
                    label_key = 'roi_labels' if 'roi_labels' in batch_dict else 'batch_pred_labels'
                    label_preds = batch_dict[label_key][index]
                else:
                    label_preds = None

                cur_start_idx = 0
                pred_scores, pred_labels, pred_boxes = [], [], []
                for cur_cls_preds, cur_label_mapping in zip(cls_preds, multihead_label_mapping):
                    # assert cur_cls_preds.shape[1] == len(cur_label_mapping)
                    cur_box_preds = box_preds[cur_start_idx: cur_start_idx + cur_cls_preds.shape[0]]
                    cur_label_preds = label_preds[cur_start_idx: cur_start_idx + cur_cls_preds.shape[0]] if label_preds is not None else None
                    cur_pred_scores, cur_pred_labels, cur_pred_boxes = model_nms_utils.multi_classes_nms(
                        cls_scores=cur_cls_preds, box_preds=cur_box_preds,
                        nms_config=post_process_cfg.NMS_CONFIG,
                        score_thresh=post_process_cfg.SCORE_THRESH,
                        label_preds=cur_label_preds
                    )
                    cur_pred_labels = cur_label_mapping[cur_pred_labels]
                    pred_scores.append(cur_pred_scores)
                    pred_labels.append(cur_pred_labels)
                    pred_boxes.append(cur_pred_boxes)
                    cur_start_idx += cur_cls_preds.shape[0]

                final_scores = torch.cat(pred_scores, dim=0)
                final_labels = torch.cat(pred_labels, dim=0)
                final_boxes = torch.cat(pred_boxes, dim=0)
            else:
                cls_preds, label_preds = torch.max(cls_preds, dim=-1)
                if batch_dict.get('has_class_labels', False):
                    label_key = 'roi_labels' if 'roi_labels' in batch_dict else 'batch_pred_labels'
                    label_preds = batch_dict[label_key][index]
                else:
                    label_preds = label_preds + 1
                selected, selected_scores = model_nms_utils.class_agnostic_nms(
                    box_scores=cls_preds, box_preds=box_preds,
                    nms_config=post_process_cfg.NMS_CONFIG,
                    score_thresh=post_process_cfg.SCORE_THRESH
                )

                if post_process_cfg.OUTPUT_RAW_SCORE:
                    max_cls_preds, _ = torch.max(src_cls_preds, dim=-1)
                    selected_scores = max_cls_preds[selected]

                final_scores = selected_scores
                final_labels = label_preds[selected]
                final_boxes = box_preds[selected]

            record_dict = {
                'pred_boxes': final_boxes,
                'pred_scores': final_scores,
                'pred_labels': final_labels,
            }
            # TODO: temporally set
            recall_dict, iou_results, ioubev_results = self.generate_recall_record_legacy(  ##############################
                box_preds=final_boxes if 'rois' not in batch_dict else src_box_preds,
                recall_dict=recall_dict, batch_index=index, data_dict=batch_dict,
                thresh_list=post_process_cfg.RECALL_THRESH_LIST
            )
            if iou_results is not None:
                record_dict['iou_results'] = iou_results
                record_dict['ioubev_results'] = ioubev_results

            pred_dicts.append(record_dict)

        return pred_dicts, recall_dict

    @staticmethod
    def generate_recall_record(box_preds, recall_dict, batch_index, data_dict=None, thresh_list=None):
        if 'gt_boxes' not in data_dict:
            return recall_dict

        rois = data_dict['rois'][batch_index] if 'rois' in data_dict else None
        gt_boxes = data_dict['gt_boxes'][batch_index]

        if recall_dict.__len__() == 0:
            recall_dict = {'gt': 0}
            for cur_thresh in thresh_list:
                recall_dict['roi_%s' % (str(cur_thresh))] = 0
                recall_dict['rcnn_%s' % (str(cur_thresh))] = 0

        cur_gt = gt_boxes
        k = cur_gt.__len__() - 1
        while k >= 0 and cur_gt[k].sum() == 0:
            k -= 1
        cur_gt = cur_gt[:k + 1]

        if cur_gt.shape[0] > 0:
            if box_preds.shape[0] > 0:
                iou3d_rcnn = iou3d_nms_utils.boxes_iou3d_gpu(box_preds[:, 0:7], cur_gt[:, 0:7])
            else:
                iou3d_rcnn = torch.zeros((0, cur_gt.shape[0]))

            if rois is not None:
                iou3d_roi = iou3d_nms_utils.boxes_iou3d_gpu(rois[:, 0:7], cur_gt[:, 0:7])

            for cur_thresh in thresh_list:
                if iou3d_rcnn.shape[0] == 0:
                    recall_dict['rcnn_%s' % str(cur_thresh)] += 0
                else:
                    rcnn_recalled = (iou3d_rcnn.max(dim=0)[0] > cur_thresh).sum().item()
                    recall_dict['rcnn_%s' % str(cur_thresh)] += rcnn_recalled
                if rois is not None:
                    roi_recalled = (iou3d_roi.max(dim=0)[0] > cur_thresh).sum().item()
                    recall_dict['roi_%s' % str(cur_thresh)] += roi_recalled

            recall_dict['gt'] += cur_gt.shape[0]
        else:
            gt_iou = box_preds.new_zeros(box_preds.shape[0])
        return recall_dict

    @staticmethod
    def generate_recall_record_legacy(box_preds, recall_dict, batch_index, data_dict=None, thresh_list=None):
        if 'gt_boxes' not in data_dict:
            return recall_dict, None, None

        rois = data_dict['rois'][batch_index] if 'rois' in data_dict else None
        gt_boxes = data_dict['gt_boxes'][batch_index]

        if recall_dict.__len__() == 0:
            recall_dict = {'gt': 0}
            for cur_thresh in thresh_list:
                recall_dict['roi_%s' % (str(cur_thresh))] = 0
                recall_dict['rcnn_%s' % (str(cur_thresh))] = 0

        cur_gt = gt_boxes
        k = cur_gt.__len__() - 1
        while k > 0 and cur_gt[k].sum() == 0:
            k -= 1
        cur_gt = cur_gt[:k + 1]

        if cur_gt.shape[0] > 0:
            if box_preds.shape[0] > 0:
                iou3d_rcnn = iou3d_nms_utils.boxes_iou3d_gpu(
                    box_preds[:, 0:7], cur_gt[:, 0:7])
                ioubev_rcnn = iou3d_nms_utils.boxes_iou_bev(
                    box_preds[:, 0:7], cur_gt[:, 0:7])
            else:
                iou3d_rcnn = torch.zeros((0, cur_gt.shape[0]))
                ioubev_rcnn = torch.zeros((0, cur_gt.shape[00]))

            if rois is not None:
                iou3d_roi = iou3d_nms_utils.boxes_iou3d_gpu(
                    rois[:, 0:7], cur_gt[:, 0:7])

            for cur_thresh in thresh_list:
                if iou3d_rcnn.shape[0] == 0:
                    recall_dict['rcnn_%s' % str(cur_thresh)] += 0
                else:
                    rcnn_recalled = (iou3d_rcnn.max(dim=0)[
                                     0] > cur_thresh).sum().item()
                    recall_dict['rcnn_%s' % str(cur_thresh)] += rcnn_recalled
                if rois is not None:
                    roi_recalled = (iou3d_roi.max(dim=0)[
                                    0] > cur_thresh).sum().item()
                    recall_dict['roi_%s' % str(cur_thresh)] += roi_recalled

            # per box iou
            if iou3d_rcnn.shape[0] == 0:
                iou_results = [0.] * cur_gt.shape[0]
                ioubev_results = [0.] * cur_gt.shape[0]
            else:
                iou_results = iou3d_rcnn.max(0).values.cpu().numpy().tolist()
                ioubev_results = ioubev_rcnn.max(0).values.cpu().numpy().tolist()

            recall_dict['gt'] += cur_gt.shape[0]
        else:
            gt_iou = box_preds.new_zeros(box_preds.shape[0])
            iou_results = []
            ioubev_results = []
        return recall_dict, iou_results, ioubev_results

    def load_params_from_file(self, filename, logger, to_cpu=False, strict=False):
        if not os.path.isfile(filename):
            return 

        logger.info('==> (Stereo) Loading parameters from checkpoint %s to %s' %
                    (filename, 'CPU' if to_cpu else 'GPU'))
        loc_type = torch.device('cpu') if to_cpu else None
        checkpoint = torch.load(filename, map_location=loc_type)
        model_state_disk = checkpoint['model_state']

        update_model_state = {}
        for key, val in model_state_disk.items():
            if key in self.state_dict() and self.state_dict()[key].shape == model_state_disk[key].shape:
                update_model_state[key] = val
            
                # logger.info('Update weight %s: %s' % (key, str(val.shape)))

        state_dict = self.state_dict()
        state_dict.update(update_model_state)
        self.load_state_dict(state_dict, strict=strict)

        for key in state_dict:
            if key not in update_model_state:
                logger.info('Not updated weight %s: %s <- %s' %
                            (key, str(state_dict[key].shape), str(model_state_disk[key].shape) if key in model_state_disk else 'None'))

        logger.info('==> Done (loaded %d/%d)' %
                    (len(update_model_state), len(self.state_dict())))

    def load_params_with_optimizer(self, filename, to_cpu=False, optimizer=None, logger=None):
        if not os.path.isfile(filename):
            raise FileNotFoundError

        logger.info('==> Loading parameters from checkpoint %s to %s' %
                    (filename, 'CPU' if to_cpu else 'GPU'))
        loc_type = torch.device('cpu') if to_cpu else None
        checkpoint = torch.load(filename, map_location=loc_type)
        epoch = checkpoint.get('epoch', -1)
        it = checkpoint.get('it', 0.0)
        # import pdb
        # pdb.set_trace()
        # checkpoint['model_state'].pop()
        # checkpoint['model_state']['dense_head_2d.bbox_head.atss_cls.bias'] = checkpoint['model_state']['dense_head_2d.bbox_head.atss_cls.bias'][7:8]
        # checkpoint['model_state']['dense_head_2d.bbox_head.atss_reg.weight'] = checkpoint['model_state']['dense_head_2d.bbox_head.atss_reg.weight'][28:32, ...]
        # checkpoint['model_state']['dense_head_2d.bbox_head.atss_reg.bias'] = checkpoint['model_state']['dense_head_2d.bbox_head.atss_reg.bias'][28:32]
        # checkpoint['model_state']['dense_head.conv_cls.weight'] = checkpoint['model_state']['dense_head.conv_cls.weight'][140:160, ...]
        # checkpoint['model_state']['dense_head.conv_cls.bias'] = checkpoint['model_state']['dense_head.conv_cls.bias'][140:160]
        # for k, v in checkpoint['model_state'].items():
        #     if k in ['dense_head_2d.bbox_head.atss_cls.bias', 'dense_head_2d.bbox_head.atss_reg.weight', 'dense_head_2d.bbox_head.atss_reg.bias',
        #              'dense_head.conv_cls.weight', 'dense_head.conv_cls.bias']:
                # del checkpoint['model_state'][k]
        # for k in ['dense_head_2d.bbox_head.atss_cls.bias', 'dense_head_2d.bbox_head.atss_reg.weight', 'dense_head_2d.bbox_head.atss_reg.bias',
        #              'dense_head.conv_cls.weight', 'dense_head.conv_cls.bias', 'dense_head_2d.bbox_head.atss_cls.weight']:
        #     del checkpoint['model_state'][k]
        
        self.load_state_dict(checkpoint['model_state'])

        if optimizer is not None:
            if 'optimizer_state' in checkpoint and checkpoint['optimizer_state'] is not None:
                logger.info('==> Loading optimizer parameters from checkpoint %s to %s'
                            % (filename, 'CPU' if to_cpu else 'GPU'))
                optimizer.load_state_dict(checkpoint['optimizer_state'])
            else:
                assert filename[-4] == '.', filename
                src_file, ext = filename[:-4], filename[-3:]
                optimizer_filename = '%s_optim.%s' % (src_file, ext)
                if os.path.exists(optimizer_filename):
                    optimizer_ckpt = torch.load(
                        optimizer_filename, map_location=loc_type)
                    optimizer.load_state_dict(
                        optimizer_ckpt['optimizer_state'])

        logger.info('==> Done')

        return it, epoch